incorporate suggestions

Author: rahul2393

samples/samples/admin/pg_samples.py

@@ -19,16 +19,74 @@
 For more information, see the README.rst under /spanner.
 """
 from google.cloud import spanner
-from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+from google.cloud.spanner_admin_database_v1.types.common import DatabaseDialect
 
 OPERATION_TIMEOUT_SECONDS = 240
 
 
+# [START spanner_postgresql_create_database]
+def create_database(instance_id, database_id):
+    """Creates a PostgreSql database and tables for sample data."""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+
+    request = spanner_database_admin.CreateDatabaseRequest(
+        parent=instance.name,
+        create_statement=f'CREATE DATABASE "{database_id}"',
+        database_dialect=DatabaseDialect.POSTGRESQL,
+    )
+
+    operation = spanner_client.database_admin_api.create_database(request=request)
+
+    print("Waiting for operation to complete...")
+    database = operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    create_table_using_ddl(database.name)
+    print("Created database {} on instance {}".format(database_id, instance_id))
+
+
+def create_table_using_ddl(database_name):
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database_name,
+        statements=[
+            """CREATE TABLE Singers (
+  SingerId   bigint NOT NULL,
+  FirstName  character varying(1024),
+  LastName   character varying(1024),
+  SingerInfo bytea,
+  FullName   character varying(2048)
+    GENERATED ALWAYS AS (FirstName || ' ' || LastName) STORED,
+  PRIMARY KEY (SingerId)
+  )""",
+            """CREATE TABLE Albums (
+  SingerId     bigint NOT NULL,
+  AlbumId      bigint NOT NULL,
+  AlbumTitle   character varying(1024),
+  PRIMARY KEY (SingerId, AlbumId)
+  ) INTERLEAVE IN PARENT Singers ON DELETE CASCADE""",
+        ],
+    )
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+
+# [END spanner_postgresql_create_database]
+
+
 def create_table_with_datatypes(instance_id, database_id):
     """Creates a table with supported datatypes."""
     # [START spanner_postgresql_create_table_with_datatypes]
     # instance_id = "your-spanner-instance"
     # database_id = "your-spanner-db-id"
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
     spanner_client = spanner.Client()
     instance = spanner_client.instance(instance_id)
     database = instance.database(database_id)
@@ -61,6 +119,31 @@ def create_table_with_datatypes(instance_id, database_id):
     # [END spanner_postgresql_create_table_with_datatypes]
 
 
+# [START spanner_postgresql_add_column]
+def add_column(instance_id, database_id):
+    """Adds a new column to the Albums table in the example database."""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database.name,
+        statements=["ALTER TABLE Albums ADD COLUMN MarketingBudget BIGINT"],
+    )
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print("Added the MarketingBudget column.")
+
+
+# [END spanner_postgresql_add_column]
+
+
 # [START spanner_postgresql_jsonb_add_column]
 def add_jsonb_column(instance_id, database_id):
     """
@@ -81,6 +164,8 @@ def add_jsonb_column(instance_id, database_id):
     # instance_id = "your-spanner-instance"
     # database_id = "your-spanner-db-id"
 
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
     spanner_client = spanner.Client()
     instance = spanner_client.instance(instance_id)
     database = instance.database(database_id)
@@ -103,3 +188,164 @@ def add_jsonb_column(instance_id, database_id):
 
 
 # [END spanner_postgresql_jsonb_add_column]
+
+
+# [START spanner_postgresql_create_storing_index]
+def add_storing_index(instance_id, database_id):
+    """Adds an storing index to the example database."""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database.name,
+        statements=[
+            "CREATE INDEX AlbumsByAlbumTitle2 ON Albums(AlbumTitle)"
+            "INCLUDE (MarketingBudget)"
+        ],
+    )
+
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print("Added the AlbumsByAlbumTitle2 index.")
+
+
+# [END spanner_postgresql_create_storing_index]
+
+
+# [START spanner_postgresql_create_sequence]
+def create_sequence(instance_id, database_id):
+    """Creates the Sequence and insert data"""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database.name,
+        statements=[
+            "CREATE SEQUENCE Seq BIT_REVERSED_POSITIVE",
+            """CREATE TABLE Customers (
+        CustomerId  BIGINT DEFAULT nextval('Seq'),
+        CustomerName  character varying(1024),
+        PRIMARY KEY (CustomerId)
+        )""",
+        ],
+    )
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print(
+        "Created Seq sequence and Customers table, where the key column CustomerId uses the sequence as a default value on database {} on instance {}".format(
+            database_id, instance_id
+        )
+    )
+
+    def insert_customers(transaction):
+        results = transaction.execute_sql(
+            "INSERT INTO Customers (CustomerName) VALUES "
+            "('Alice'), "
+            "('David'), "
+            "('Marc') "
+            "RETURNING CustomerId"
+        )
+        for result in results:
+            print("Inserted customer record with Customer Id: {}".format(*result))
+        print(
+            "Number of customer records inserted is {}".format(
+                results.stats.row_count_exact
+            )
+        )
+
+    database.run_in_transaction(insert_customers)
+
+
+# [END spanner_postgresql_create_sequence]
+
+
+# [START spanner_postgresql_alter_sequence]
+def alter_sequence(instance_id, database_id):
+    """Alters the Sequence and insert data"""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database.name,
+        statements=["ALTER SEQUENCE Seq SKIP RANGE 1000 5000000"],
+    )
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print(
+        "Altered Seq sequence to skip an inclusive range between 1000 and 5000000 on database {} on instance {}".format(
+            database_id, instance_id
+        )
+    )
+
+    def insert_customers(transaction):
+        results = transaction.execute_sql(
+            "INSERT INTO Customers (CustomerName) VALUES "
+            "('Lea'), "
+            "('Cataline'), "
+            "('Smith') "
+            "RETURNING CustomerId"
+        )
+        for result in results:
+            print("Inserted customer record with Customer Id: {}".format(*result))
+        print(
+            "Number of customer records inserted is {}".format(
+                results.stats.row_count_exact
+            )
+        )
+
+    database.run_in_transaction(insert_customers)
+
+
+# [END spanner_postgresql_alter_sequence]
+
+
+# [START spanner_postgresql_drop_sequence]
+def drop_sequence(instance_id, database_id):
+    """Drops the Sequence"""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database.name,
+        statements=[
+            "ALTER TABLE Customers ALTER COLUMN CustomerId DROP DEFAULT",
+            "DROP SEQUENCE Seq",
+        ],
+    )
+    operation = spanner_client.database_admin_api.update_database_ddl(request)
+
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print(
+        "Altered Customers table to drop DEFAULT from CustomerId column and dropped the Seq sequence on database {} on instance {}".format(
+            database_id, instance_id
+        )
+    )
+
+
+# [END spanner_postgresql_drop_sequence]

samples/samples/admin/pg_samples_test.py

@@ -104,16 +104,75 @@ def default_leader():
     return "us-east4"
 
 
+@pytest.mark.dependency(name="create_database")
+def test_create_database_explicit(sample_instance, create_database_id):
+    # Rather than re-use 'sample_database', we create a new database, to
+    # ensure that the 'create_database' snippet is tested.
+    samples.create_database(sample_instance.instance_id, create_database_id)
+    database = sample_instance.database(create_database_id)
+    database.drop()
+
+
 @pytest.mark.dependency(name="create_table_with_datatypes")
 def test_create_table_with_datatypes(capsys, instance_id, sample_database):
     samples.create_table_with_datatypes(instance_id, sample_database.database_id)
     out, _ = capsys.readouterr()
     assert "Created Venues table on database" in out
 
 
-@pytest.mark.dependency(name="add_jsonb_column", depends=["insert_datatypes_data"])
+@pytest.mark.dependency(name="add_column", depends=["create_database"])
+def test_add_column(capsys, instance_id, sample_database):
+    samples.add_column(instance_id, sample_database.database_id)
+    out, _ = capsys.readouterr()
+    assert "Added the MarketingBudget column." in out
+
+
+@pytest.mark.dependency(name="add_storing_index", depends=["create_database"])
+def test_add_storing_index(capsys, instance_id, sample_database):
+    samples.add_storing_index(instance_id, sample_database.database_id)
+    out, _ = capsys.readouterr()
+    assert "Added the AlbumsByAlbumTitle2 index." in out
+
+
+@pytest.mark.dependency(
+    name="add_jsonb_column", depends=["create_table_with_datatypes"]
+)
 def test_add_jsonb_column(capsys, instance_id, sample_database):
     samples.add_jsonb_column(instance_id, sample_database.database_id)
     out, _ = capsys.readouterr()
     assert "Waiting for operation to complete..." in out
     assert 'Altered table "Venues" on database ' in out
+
+
+@pytest.mark.dependency(name="create_sequence")
+def test_create_sequence(capsys, instance_id, bit_reverse_sequence_database):
+    samples.create_sequence(instance_id, bit_reverse_sequence_database.database_id)
+    out, _ = capsys.readouterr()
+    assert (
+        "Created Seq sequence and Customers table, where the key column CustomerId uses the sequence as a default value on database"
+        in out
+    )
+    assert "Number of customer records inserted is 3" in out
+    assert "Inserted customer record with Customer Id:" in out
+
+
+@pytest.mark.dependency(name="alter_sequence", depends=["create_sequence"])
+def test_alter_sequence(capsys, instance_id, bit_reverse_sequence_database):
+    samples.alter_sequence(instance_id, bit_reverse_sequence_database.database_id)
+    out, _ = capsys.readouterr()
+    assert (
+        "Altered Seq sequence to skip an inclusive range between 1000 and 5000000 on database"
+        in out
+    )
+    assert "Number of customer records inserted is 3" in out
+    assert "Inserted customer record with Customer Id:" in out
+
+
+@pytest.mark.dependency(depends=["alter_sequence"])
+def test_drop_sequence(capsys, instance_id, bit_reverse_sequence_database):
+    samples.drop_sequence(instance_id, bit_reverse_sequence_database.database_id)
+    out, _ = capsys.readouterr()
+    assert (
+        "Altered Customers table to drop DEFAULT from CustomerId column and dropped the Seq sequence on database"
+        in out
+    )