retrieval_service/datastore/providers/spanner_gsql.py [47:499]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - project: str instance: str database: str service_account_key_file: Optional[str] = None # Client class for interacting with Spanner class Client(datastore.Client[Config]): OPERATION_TIMEOUT_SECONDS = 240 BATCH_SIZE = 1000 AIRPORT_COLUMNS = ["id", "iata", "name", "city", "country"] AMENITIES_COLUMNS = [ "id", "name", "description", "location", "terminal", "category", "hour", "sunday_start_hour", "sunday_end_hour", "monday_start_hour", "monday_end_hour", "tuesday_start_hour", "tuesday_end_hour", "wednesday_start_hour", "wednesday_end_hour", "thursday_start_hour", "thursday_end_hour", "friday_start_hour", "friday_end_hour", "saturday_start_hour", "saturday_end_hour", "content", "embedding", ] FLIGHTS_COLUMNS = [ "id", "airline", "flight_number", "departure_airport", "arrival_airport", "departure_time", "arrival_time", "departure_gate", "arrival_gate", ] POLICIES_COLUMNS = ["id", "content", "embedding"] """ Client class for interacting with Spanner. Attributes: __client (spanner.Client): Spanner client instance. __instance_id (str): ID of the Spanner instance. __database_id (str): ID of the Spanner database. __instance (Instance): Spanner instance. __database (Database): Spanner database. """ @datastore.classproperty def kind(cls): return SPANNER_IDENTIFIER def __init__(self, client: spanner.Client, instance_id: str, database_id: str): """ Initialize the Spanner client. Args: client (spanner.Client): Spanner client instance. instance_id (str): ID of the Spanner instance. database_id (str): ID of the Spanner database. """ self.__client = client self.__instance_id = instance_id self.__database_id = database_id self.__instance = self.__client.instance(self.__instance_id) self.__database = self.__instance.database(self.__database_id) @classmethod async def create(cls, config: Config) -> "Client": """ Create a Spanner client. Args: config (Config): Configuration for creating the client. Returns: Client: Initialized Spanner client. """ client: spanner.Client if config.service_account_key_file is not None: credentials = service_account.Credentials.from_service_account_file( config.service_account_key_file ) client = spanner.Client(project=config.project, credentials=credentials) else: client = spanner.Client(project=config.project) instance_id = config.instance instance = client.instance(instance_id) if not instance.exists(): raise Exception(f"Instance with id: {instance_id} doesn't exist.") database_id = config.database database = instance.database(database_id) if not database.exists(): raise Exception(f"Database with id: {database_id} doesn't exist.") return cls(client, instance_id, database_id) async def initialize_data( self, airports: list[models.Airport], amenities: list[models.Amenity], flights: list[models.Flight], policies: list[models.Policy], ) -> None: """ Initialize data in the Spanner database by creating tables and inserting records. Args: airports (list[models.Airport]): list of airports to be initialized. amenities (list[models.Amenity]): list of amenities to be initialized. flights (list[models.Flight]): list of flights to be initialized. policies (list[models.Policy]): list of policies to be initialized. Returns: None """ # Initialize a list to store Data Definition Language (DDL) statements ddl = [] # Create DDL statement to drop the 'airports' table if it exists ddl.append("DROP TABLE IF EXISTS airports") # Create DDL statement to create the 'airports' table ddl.append( """ CREATE TABLE airports( id INT64, iata STRING(MAX), name STRING(MAX), city STRING(MAX), country STRING(MAX) ) PRIMARY KEY(id) """ ) # Create DDL statement to drop the 'amenities' table if it exists ddl.append("DROP TABLE IF EXISTS amenities") # Create DDL statement to create the 'amenities' table ddl.append( """ CREATE TABLE amenities( id INT64, name STRING(MAX), description STRING(MAX), location STRING(MAX), terminal STRING(MAX), category STRING(MAX), hour STRING(MAX), sunday_start_hour STRING(100), sunday_end_hour STRING(100), monday_start_hour STRING(100), monday_end_hour STRING(100), tuesday_start_hour STRING(100), tuesday_end_hour STRING(100), wednesday_start_hour STRING(100), wednesday_end_hour STRING(100), thursday_start_hour STRING(100), thursday_end_hour STRING(100), friday_start_hour STRING(100), friday_end_hour STRING(100), saturday_start_hour STRING(100), saturday_end_hour STRING(100), content STRING(MAX) NOT NULL, embedding ARRAY NOT NULL ) PRIMARY KEY(id) """ ) # Create DDL statement to drop the 'flights' table if it exists ddl.append("DROP TABLE IF EXISTS flights") # Create DDL statement to create the 'flights' table ddl.append( """ CREATE TABLE flights( id INT64, airline STRING(MAX), flight_number STRING(MAX), departure_airport STRING(MAX), arrival_airport STRING(MAX), departure_time STRING(100), arrival_time STRING(100), departure_gate STRING(MAX), arrival_gate STRING(MAX) ) PRIMARY KEY(id) """ ) # Create DDL statement to drop the 'policies' table if it exists ddl.append("DROP TABLE IF EXISTS policies") # Create DDL statement to create the 'policies' table ddl.append( """ CREATE TABLE policies( id INT64, content STRING(MAX) NOT NULL, embedding ARRAY NOT NULL ) PRIMARY KEY(id) """ ) # Create DDL statement to drop the 'tickets' table if it exists ddl.append("DROP TABLE IF EXISTS tickets") # Create DDL statement to create the 'tickets' table ddl.append( """ CREATE TABLE tickets( user_id STRING(MAX), user_name STRING(MAX), user_email STRING(MAX), airline STRING(MAX), flight_number STRING(MAX), departure_airport STRING(MAX), arrival_airport STRING(MAX), departure_time STRING(100), arrival_time STRING(100) ) PRIMARY KEY(user_id, airline, flight_number, departure_time) """ ) # Update the schema using DDL statements operation = self.__database.update_ddl(ddl) print("Waiting for schema update operation to complete...") operation.result(self.OPERATION_TIMEOUT_SECONDS) print("Schema update operation completed") # Insert data into 'airports' table using batch operation values = [ tuple(getattr(airport, field) for field in self.AIRPORT_COLUMNS) for airport in airports ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="airports", columns=self.AIRPORT_COLUMNS, values=records, ) # Insert data into 'amenities' table using batch operation values = [ tuple( ( str(getattr(amenity, field)) if isinstance(getattr(amenity, field), datetime.time) else getattr(amenity, field) ) for field in self.AMENITIES_COLUMNS ) for amenity in amenities ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="amenities", columns=self.AMENITIES_COLUMNS, values=records, ) # Insert data into 'flights' table using batch operation values = [ tuple( ( str(getattr(flight, field)) if isinstance(getattr(flight, field), datetime.datetime) else getattr(flight, field) ) for field in self.FLIGHTS_COLUMNS ) for flight in flights ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="flights", columns=self.FLIGHTS_COLUMNS, values=records, ) # Insert data into 'policies' table using batch operation values = [ tuple(getattr(policy, field) for field in self.POLICIES_COLUMNS) for policy in policies ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="policies", columns=self.POLICIES_COLUMNS, values=records, ) # Return None to indicate successful initialization return None async def export_data( self, ) -> tuple[ list[models.Airport], list[models.Amenity], list[models.Flight], list[models.Policy], ]: """ Export data from the Spanner database. Returns: tuple: A tuple containing lists of airports, amenities, flights, and policies. """ airports: list = [] amenities: list = [] flights: list = [] policies: list = [] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables airport_results = snapshot.execute_sql( "SELECT {} FROM airports ORDER BY id ASC".format( ",".join(self.AIRPORT_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch airports: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method airports = [ models.Airport.model_validate( {key: value for key, value in zip(self.AIRPORT_COLUMNS, a)} ) for a in airport_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables amenity_results = snapshot.execute_sql( "SELECT {} FROM amenities ORDER BY id ASC".format( ",".join(self.AMENITIES_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch amenities: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method amenities = [ models.Amenity.model_validate( {key: value for key, value in zip(self.AMENITIES_COLUMNS, a)} ) for a in amenity_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables flights_results = snapshot.execute_sql( "SELECT {} FROM flights ORDER BY id ASC".format( ",".join(self.FLIGHTS_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch flights: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method flights = [ models.Flight.model_validate( {key: value for key, value in zip(self.FLIGHTS_COLUMNS, a)} ) for a in flights_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables policy_results = snapshot.execute_sql( "SELECT {} FROM policies ORDER BY id ASC".format( ",".join(self.POLICIES_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch policies: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method policies = [ models.Policy.model_validate( {key: value for key, value in zip(self.POLICIES_COLUMNS, a)} ) for a in policy_results ] return airports, amenities, flights, policies async def get_airport_by_id( self, id: int ) -> tuple[Optional[models.Airport], Optional[str]]: """ Retrieve an airport by its ID. Args: id (int): The ID of the airport. Returns: Optional[models.Airport]: An Airport model instance if found, else None. """ with self.__database.snapshot() as snapshot: # Execute SQL query to fetch airport by ID result = snapshot.execute_sql( - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - retrieval_service/datastore/providers/spanner_postgres.py [47:500]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - project: str instance: str database: str service_account_key_file: Optional[str] = None # Client class for interacting with Spanner class Client(datastore.Client[Config]): OPERATION_TIMEOUT_SECONDS = 240 BATCH_SIZE = 1000 AIRPORT_COLUMNS = ["id", "iata", "name", "city", "country"] AMENITIES_COLUMNS = [ "id", "name", "description", "location", "terminal", "category", "hour", "sunday_start_hour", "sunday_end_hour", "monday_start_hour", "monday_end_hour", "tuesday_start_hour", "tuesday_end_hour", "wednesday_start_hour", "wednesday_end_hour", "thursday_start_hour", "thursday_end_hour", "friday_start_hour", "friday_end_hour", "saturday_start_hour", "saturday_end_hour", "content", "embedding", ] FLIGHTS_COLUMNS = [ "id", "airline", "flight_number", "departure_airport", "arrival_airport", "departure_time", "arrival_time", "departure_gate", "arrival_gate", ] POLICIES_COLUMNS = ["id", "content", "embedding"] """ Client class for interacting with Spanner. Attributes: __client (spanner.Client): Spanner client instance. __instance_id (str): ID of the Spanner instance. __database_id (str): ID of the Spanner database. __instance (Instance): Spanner instance. __database (Database): Spanner database. """ @datastore.classproperty def kind(cls): return SPANNER_IDENTIFIER def __init__(self, client: spanner.Client, instance_id: str, database_id: str): """ Initialize the Spanner client. Args: client (spanner.Client): Spanner client instance. instance_id (str): ID of the Spanner instance. database_id (str): ID of the Spanner database. """ self.__client = client self.__instance_id = instance_id self.__database_id = database_id self.__instance = self.__client.instance(self.__instance_id) self.__database = self.__instance.database(self.__database_id) @classmethod async def create(cls, config: Config) -> "Client": """ Create a Spanner client. Args: config (Config): Configuration for creating the client. Returns: Client: Initialized Spanner client. """ client: spanner.Client if config.service_account_key_file is not None: credentials = service_account.Credentials.from_service_account_file( config.service_account_key_file ) client = spanner.Client(project=config.project, credentials=credentials) else: client = spanner.Client(project=config.project) instance_id = config.instance instance = client.instance(instance_id) if not instance.exists(): raise Exception(f"Instance with id: {instance_id} doesn't exist.") database_id = config.database database = instance.database(database_id) if not database.exists(): raise Exception(f"Database with id: {database_id} doesn't exist.") return cls(client, instance_id, database_id) async def initialize_data( self, airports: list[models.Airport], amenities: list[models.Amenity], flights: list[models.Flight], policies: list[models.Policy], ) -> None: """ Initialize data in the Spanner database by creating tables and inserting records. Args: airports (list[models.Airport]): list of airports to be initialized. amenities (list[models.Amenity]): list of amenities to be initialized. flights (list[models.Flight]): list of flights to be initialized. policies (list[models.Policy]): list of policies to be initialized. Returns: None """ # Initialize a list to store Data Definition Language (DDL) statements ddl = [] # Create DDL statement to drop the 'airports' table if it exists ddl.append("DROP TABLE IF EXISTS airports") # Create DDL statement to create the 'airports' table ddl.append( """ CREATE TABLE airports( id BIGINT PRIMARY KEY, iata VARCHAR, name VARCHAR, city VARCHAR, country VARCHAR ) """ ) # Create DDL statement to drop the 'amenities' table if it exists ddl.append("DROP TABLE IF EXISTS amenities") # Create DDL statement to create the 'amenities' table ddl.append( """ CREATE TABLE amenities( id BIGINT PRIMARY KEY, name VARCHAR, description VARCHAR, location VARCHAR, terminal VARCHAR, category VARCHAR, hour VARCHAR, sunday_start_hour VARCHAR, sunday_end_hour VARCHAR, monday_start_hour VARCHAR, monday_end_hour VARCHAR, tuesday_start_hour VARCHAR, tuesday_end_hour VARCHAR, wednesday_start_hour VARCHAR, wednesday_end_hour VARCHAR, thursday_start_hour VARCHAR, thursday_end_hour VARCHAR, friday_start_hour VARCHAR, friday_end_hour VARCHAR, saturday_start_hour VARCHAR, saturday_end_hour VARCHAR, content VARCHAR NOT NULL, embedding FLOAT8[] NOT NULL ) """ ) # Create DDL statement to drop the 'flights' table if it exists ddl.append("DROP TABLE IF EXISTS flights") # Create DDL statement to create the 'flights' table ddl.append( """ CREATE TABLE flights( id BIGINT PRIMARY KEY, airline VARCHAR, flight_number VARCHAR, departure_airport VARCHAR, arrival_airport VARCHAR, departure_time VARCHAR(100), arrival_time VARCHAR(100), departure_gate VARCHAR, arrival_gate VARCHAR ) """ ) # Create DDL statement to drop the 'policies' table if it exists ddl.append("DROP TABLE IF EXISTS policies") # Create DDL statement to create the 'policies' table ddl.append( """ CREATE TABLE policies( id BIGINT PRIMARY KEY, content VARCHAR NOT NULL, embedding FLOAT8[] NOT NULL ) """ ) # Create DDL statement to drop the 'tickets' table if it exists ddl.append("DROP TABLE IF EXISTS tickets") # Create DDL statement to create the 'tickets' table ddl.append( """ CREATE TABLE tickets( user_id VARCHAR, user_name VARCHAR, user_email VARCHAR, airline VARCHAR, flight_number VARCHAR, departure_airport VARCHAR, arrival_airport VARCHAR, departure_time VARCHAR(100), arrival_time VARCHAR(100), PRIMARY KEY(user_id, airline, flight_number, departure_time) ) """ ) # Update the schema using DDL statements operation = self.__database.update_ddl(ddl) print("Waiting for schema update operation to complete...") operation.result(self.OPERATION_TIMEOUT_SECONDS) print("Schema update operation completed") # Insert data into 'airports' table using batch operation values = [ tuple(getattr(airport, field) for field in self.AIRPORT_COLUMNS) for airport in airports ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="airports", columns=self.AIRPORT_COLUMNS, values=records, ) # Insert data into 'amenities' table using batch operation values = [ tuple( ( str(getattr(amenity, field)) if isinstance(getattr(amenity, field), datetime.time) else getattr(amenity, field) ) for field in self.AMENITIES_COLUMNS ) for amenity in amenities ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="amenities", columns=self.AMENITIES_COLUMNS, values=records, ) # Insert data into 'flights' table using batch operation values = [ tuple( ( str(getattr(flight, field)) if isinstance(getattr(flight, field), datetime.datetime) else getattr(flight, field) ) for field in self.FLIGHTS_COLUMNS ) for flight in flights ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="flights", columns=self.FLIGHTS_COLUMNS, values=records, ) # Insert data into 'policies' table using batch operation values = [ tuple(getattr(policy, field) for field in self.POLICIES_COLUMNS) for policy in policies ] for i in range(0, len(values), self.BATCH_SIZE): records = values[i : i + self.BATCH_SIZE] with self.__database.batch() as batch: batch.insert( table="policies", columns=self.POLICIES_COLUMNS, values=records, ) # Return None to indicate successful initialization return None async def export_data( self, ) -> tuple[ list[models.Airport], list[models.Amenity], list[models.Flight], list[models.Policy], ]: """ Export data from the Spanner database. Returns: tuple: A tuple containing lists of airports, amenities, flights, and policies. """ airports: list = [] amenities: list = [] flights: list = [] policies: list = [] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables airport_results = snapshot.execute_sql( "SELECT {} FROM airports ORDER BY id ASC".format( ",".join(self.AIRPORT_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch airports: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method airports = [ models.Airport.model_validate( {key: value for key, value in zip(self.AIRPORT_COLUMNS, a)} ) for a in airport_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables amenity_results = snapshot.execute_sql( "SELECT {} FROM amenities ORDER BY id ASC".format( ",".join(self.AMENITIES_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch amenities: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method amenities = [ models.Amenity.model_validate( {key: value for key, value in zip(self.AMENITIES_COLUMNS, a)} ) for a in amenity_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables flights_results = snapshot.execute_sql( "SELECT {} FROM flights ORDER BY id ASC".format( ",".join(self.FLIGHTS_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch flights: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method flights = [ models.Flight.model_validate( {key: value for key, value in zip(self.FLIGHTS_COLUMNS, a)} ) for a in flights_results ] try: with self.__database.snapshot() as snapshot: # Execute SQL queries to fetch data from respective tables policy_results = snapshot.execute_sql( "SELECT {} FROM policies ORDER BY id ASC".format( ",".join(self.POLICIES_COLUMNS) ) ) except Exception as e: # Handle any exceptions, such as database connection errors print(f"Error occurred while fetch policies: {e}") # Return empty lists in case of error return airports, amenities, flights, policies # Convert query results to model instances using model_validate method policies = [ models.Policy.model_validate( {key: value for key, value in zip(self.POLICIES_COLUMNS, a)} ) for a in policy_results ] return airports, amenities, flights, policies async def get_airport_by_id( self, id: int ) -> tuple[Optional[models.Airport], Optional[str]]: """ Retrieve an airport by its ID. Args: id (int): The ID of the airport. Returns: Optional[models.Airport]: An Airport model instance if found, else None. """ with self.__database.snapshot() as snapshot: # Execute SQL query to fetch airport by ID result = snapshot.execute_sql( - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -